# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# All rights reserved.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

import importlib
import os
import sys
from pathlib import Path

from transformers import AutoModelForCausalLM
from openmind.flow.model import get_tokenizer
from openmind.utils import get_logger
from openmind.flow.model.loader import get_model
from openmind.flow.arguments import get_args, initialize_openmind

logger = get_logger()


def lora_merge_and_unload(base_model: AutoModelForCausalLM, adapter_name_or_path: str):
    peft_module = importlib.import_module("peft")
    model = peft_module.PeftModel.from_pretrained(base_model, adapter_name_or_path)
    model = model.merge_and_unload()
    return model


def parse_input_to_list(input_string):
    input_string = input_string.strip()

    if input_string.startswith("[") and input_string.endswith("]"):
        input_string = input_string.replace("[", "").replace("]", "").strip()
    return input_string.split(",")


def run_export():

    if len(sys.argv) == 3 and sys.argv[-1].endswith("yaml"):
        yaml_file = sys.argv[-1]
        initialize_openmind(yaml_file)
    else:
        initialize_openmind()

    args = get_args()

    if args.adapter_name_or_path is None or args.output_dir is None:
        raise ValueError("Please set adapter_name_or_path and output path to start merge and export.")

    # create output_dir if it's not existed
    try:
        path_obj = Path(args.output_dir)
    except Exception as e:
        raise ValueError(f"The path {args.output_dir} is not a valid path format.") from e

    if not path_obj.exists():
        path_obj.mkdir(parents=True, exist_ok=True)
        logger.info_rank0(f"Directory {args.output_dir} is created for saving merge model")
    else:
        if any(os.scandir(args.output_dir)):
            raise ValueError(
                f"The output dir {args.output_dir} is not empty."
                "Please select a empty folder avoiding cover existing model files."
                "Or setting the output_dir param with a new folder name, the process will build it for you."
            )
        logger.info_rank0(f"Directory {args.output_dir} is used for saving merge model")

    model = get_model()
    tokenizer = get_tokenizer()

    # if not set per_shard_size, the max_shar_size will default to 5GB
    max_shard_size = f"{args.per_shard_size}GB" if args.per_shard_size else "5GB"

    model.save_pretrained(save_directory=args.output_dir, max_shard_size=max_shard_size, safe_serialization=True)
    tokenizer.save_pretrained(args.output_dir)
    logger.info_rank0(f"Lora weight has been successfully merged to base model and saved in {args.output_dir}")
